import xlrd
from matplotlib import pyplot as plt

from src.utils import sum_matrix_col

cate1 = ["制造", "能源", "金融", "科技", "咨询", "国防", '零售', '电商']
cate2 = ["航空", "航天", "国防", "汽车", "电脑", "软件"]

f = xlrd.open_workbook("../2018-500.xlsx")
st = f.sheet_by_index(0)

c1Idx = 0
c2Idx = 1

c1 = cate1[c1Idx]
c2 = cate2[c2Idx]

cate = c1 + " + " + c2

arr = []
for rowIndex in range(1, st.nrows):
    row = st.row(rowIndex)
    if c1 not in row[6].value or c2 not in row[6].value:
        continue
    arr.append([row[0].value, row[1].value, row[2].value, row[3].value, row[4].value, row[5].value])

dict = {}
for a in arr:
    dict.setdefault(a[5], []).append(a)

keys = []
analyze = {}
for k in dict.keys():
    rows = dict[k]
    keys.append(k)
    analyze.setdefault(k, [
        len(rows),
        sum_matrix_col(rows, 3),
        sum_matrix_col(rows, 4),
        sum_matrix_col(rows, 1, 500) - sum_matrix_col(rows, 0)
    ])

plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']

fg = plt.figure()
max5 = []
for key in keys:
    row = analyze[key]
    if len(max5) < 5:
        max5.append([key, row])
    else:
        for v in range(0, 4):
            if max5[v][1][1] < row[1]:
                max5[v] = [key, row]

X = []
Y1 = []
Y2 = []
Y3 = []
Y4 = []
for v in max5:
    X.append(v[0])
    Y1.append(v[1][0])
    Y2.append(v[1][1])
    Y3.append(v[1][2])
    Y4.append(v[1][3])

barWidth=0.5
plt.subplot(221)
plt.bar(X, Y1, barWidth)
plt.title(cate + " 企业总数")

plt.subplot(222)
plt.bar(X, Y4, barWidth)
plt.title(cate + " 企业排名变化")

plt.subplot(223)
plt.bar(X, Y2, barWidth)
plt.title(cate + " 企业收入总数")

plt.subplot(224)
plt.bar(X, Y3, barWidth)
plt.title(cate + " 企业利润总数")

# plt.suptitle("2019年世界500强 情况统计")

plt.show()
print(analyze)
